# Copyright (c) 2011, Alex Krizhevsky (akrizhevsky@gmail.com)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# - Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from data import *
import numpy as np
import numpy.random as nr


class CIFARDataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1,
                 init_batchnum=None, dp_params={}, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)
        self.data_mean = self.batch_meta['data_mean']
        self.num_colors = 3
        self.img_size = 32
        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        for d in self.data_dic:
            # This converts the data matrix to single precision and makes sure that it is C-ordered
            d['data'] = np.require((d['data'] - self.data_mean),
                                   dtype=np.single, requirements='C')
            d['labels'] = np.require(d['labels'].reshape((1, d['data'].shape[1])),
                                     dtype=np.single, requirements='C')

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
        """
            TUND
            data_dic: type(data_dic) -> dict (dictionary)
            [datadic['data'], datadic['labels']]: a list. Thats why dataIdx=0 for data and
                                                  dataIdx=1 for labels
        """
        return epoch, batchnum, [datadic['data'], datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        # TUND: 1 means single label (0, 1, 2, ...)
        return self.img_size**2 * self.num_colors if idx == 0 else 1

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3,
                                                            self.img_size, self.img_size).
                          swapaxes(1, 3).swapaxes(1, 2) / 255.0,
                          dtype=np.single)


class CroppedCIFARDataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1, init_batchnum=None,
                 dp_params=None, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['labels'] = np.require(np.tile(d['labels'].reshape((1, d['data'].shape[1])),
                                    (1, self.data_mult)), requirements='C')

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                       self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size +
                           self.inner_size].reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)

        cropped = self.cropped_data[self.batches_generated % 2]

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean
        self.batches_generated += 1
        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 1

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3,
                                                            self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                           x.shape[1]))
                    target[:, (self.num_views/2 + i) * x.shape[1]:(self.num_views/2 +i+1) * x.shape[1]] = pic[:, :, ::-1, :].reshape((self.get_data_dims(), x.shape[1]))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                (startY, startX) = (nr.randint(0, self.border_size*2 + 1),
                                    nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class TUNDCroppedCIFARDataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1, init_batchnum=None,
                 dp_params=None, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['labels'] = np.require(np.tile(d['labels'].reshape((1, d['data'].shape[1])),
                                    (1, self.data_mult)), requirements='C')

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size +
                           self.inner_size].reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)

        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1] and batchnum != 5:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 1

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3,
                                                            self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                           x.shape[1]))
                    target[:, (self.num_views/2 + i) * x.shape[1]:(self.num_views/2 +i+1) * x.shape[1]] = pic[:,:,::-1,:].reshape((self.get_data_dims(),x.shape[1]))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                (startY, startX) = (nr.randint(0, self.border_size*2 + 1),
                                    nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class DummyConvNetDataProvider(LabeledDummyDataProvider):
    def __init__(self, data_dim):
        LabeledDummyDataProvider.__init__(self, data_dim)

    def get_next_batch(self):
        epoch, batchnum, dic = LabeledDummyDataProvider.get_next_batch(self)

        dic['data'] = np.require(dic['data'].T, requirements='C')
        dic['labels'] = np.require(dic['labels'].T, requirements='C')

        return epoch, batchnum, [dic['data'], dic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    def get_data_dims(self, idx=0):
        return self.batch_meta['num_vis'] if idx == 0 else 1


class MNISTDataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1, init_batchnum=None,
                 dp_params={}, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)
        self.num_colors = 1
        self.img_size = 28
        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        for d in self.data_dic:
            # This converts the data matrix to single precision and
            # makes sure that it is C-ordered
            d['data'] = np.require(d['data'], dtype=np.single, requirements='C')
            d['labels'] = np.require(d['labels'].reshape((1, d['data'].shape[1])),
                                    dtype=np.single, requirements='C')

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
        return epoch, batchnum, [datadic['data'], datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        return self.img_size**2 * self.num_colors if idx == 0 else 1

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require(data.T.reshape(data.shape[1], 3, self.img_size,
                                        self.img_size).swapaxes(1, 3).swapaxes(1, 2)
                         / 255.0, dtype=np.single)


class KaggleGalaxyZoo256DataProvider(LabeledDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1,
                 init_batchnum=None, dp_params={}, test=False):
        LabeledDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                     init_batchnum, dp_params, test)
        self.data_mean = self.batch_meta['data_mean']
        self.num_colors = 3
        self.img_size = 256

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledDataProvider.get_next_batch(self)
        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        datadic['data'] = np.require((datadic['data'] - self.data_mean),
                                     dtype=np.single, requirements='C')
        datadic['labels'] = np.require(datadic['labels'], dtype=np.single, requirements='C')
        return epoch, batchnum, [datadic['data'], datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        # A result has 37 fields
        return self.img_size**2 * self.num_colors if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3, self.img_size,
                                                            self.img_size).swapaxes(1, 3)
                                                 .swapaxes(1, 2) / 255.0, dtype=np.single)


class KaggleGalaxyZoo32DataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1,
                 init_batchnum=None, dp_params={}, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)
        self.data_mean = self.batch_meta['data_mean']
        self.num_colors = 3
        self.img_size = 32
        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        for d in self.data_dic:
            # This converts the data matrix to single precision and makes sure that it is C-ordered
            d['data'] = np.require((d['data'] - self.data_mean),
                                   dtype=np.single, requirements='C')
            d['labels'] = np.require(d['labels'],
                                     dtype=np.single, requirements='C')

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
        return epoch, batchnum, [datadic['data'], datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        # A result has 37 fields
        return self.img_size**2 * self.num_colors if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3, self.img_size,
                                                            self.img_size).swapaxes(1, 3)
                                                 .swapaxes(1, 2) / 255.0, dtype=np.single)


class KaggleGalaxyZoo32DataProviderG(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1,
                 init_batchnum=None, dp_params={}, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)
        self.data_mean = self.batch_meta['data_mean']
        self.num_colors = 1
        self.img_size = 32
        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        for d in self.data_dic:
            # This converts the data matrix to single precision and makes sure that it is C-ordered
            d['data'] = np.require((d['data'] - self.data_mean),
                                  dtype=np.single, requirements='C')
            d['labels'] = np.require(d['labels'],
                                    dtype=np.single, requirements='C')

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
        return epoch, batchnum, [datadic['data'], datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        # A result has 37 fields
        return self.img_size**2 * self.num_colors if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], self.img_size,
                                                            self.img_size), dtype=np.single)


class CroppedKaggleGalaxyZoo256DataProvider(LabeledDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                     init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 256 - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.batch_meta['num_cases_per_batch']*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 256, 256))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledDataProvider.get_next_batch(self)
        datadic['data'] = np.require(datadic['data'], requirements='C')
        datadic['labels'] = np.require(np.tile(datadic['labels'], (1, self.data_mult)),
                                       requirements='C')

        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 256, 256, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class CroppedKaggleGalaxyZoo32DataProvider(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['labels'] = np.require(np.tile(d['labels'], (1, self.data_mult)),
                                     requirements='C')

        # cropped_data is a list with 2 elements: cropped_data[0] & cropped_data[1].
        # They have the same dimensionalities. So why do we need 2? See the answer right below
        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)

        # Here we cropped_data[0] and cropped_data[1] round-robin to avoid overlapping
        # assignments of 2 consecutive batches. Note that list assignment in Python is passing
        # reference instead of copying.
        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])
        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class CroppedKaggleGalaxyZoo32DataProviderV2(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = True
        self.num_views = 5*2
        self.data_mult = self.num_views
        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['labels'] = np.require(np.tile(d['labels'], (1, self.data_mult)),
                                     requirements='C')

        # cropped_data is a list with 2 elements: cropped_data[0] & cropped_data[1].
        # They have the same dimensionalities. So why do we need 2? See the answer right below
        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)

        # Here we cropped_data[0] and cropped_data[1] round-robin to avoid overlapping
        # assignments of 2 consecutive batches. Note that list assignment in Python is passing
        # reference instead of copying.
        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])
        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for i in xrange(self.num_views/2):
                for c in xrange(x.shape[1]):    # loop over cases
                    startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                      nr.randint(0, self.border_size*2 + 1))
                    endY, endX = startY + self.inner_size, startX + self.inner_size
                    pic = y[:, startY:endY, startX:endX, c]
                    target[:, i*x.shape[1]+c] = pic.reshape((self.get_data_dims(),))
                    pic = pic[:, :, ::-1]
                    target[:, (i+self.num_views/2)*x.shape[1]+c] = (pic.reshape(
                                                                    (self.get_data_dims(),)))


# 32x32 cropped image with horizontal reflections and rotations
# train: x2
# test: x25
class CroppedKaggleGalaxyZoo32DataProviderV3(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = test
        if test:
            self.num_trans = 5      # number of transformations
            self.num_views = 5*self.num_trans
            self.data_mult = self.num_views if self.multiview else 1
        else:
            self.num_trans = 2
            self.num_views = 1*self.num_trans
            self.data_mult = self.num_views

        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['labels'] = np.require(np.tile(d['labels'], (1, self.data_mult)),
                                     requirements='C')

        # cropped_data is a list with 2 elements: cropped_data[0] & cropped_data[1].
        # They have the same dimensionalities. So why do we need 2? See the answer right below
        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)

        # Here we cropped_data[0] and cropped_data[1] round-robin to avoid overlapping
        # assignments of 2 consecutive batches. Note that list assignment in Python is passing
        # reference instead of copying.
        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])
        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/self.num_trans):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    # pic.shape = (3, 28, 28, 10000)
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    # horizontal reflection
                    target[:, (self.num_views/self.num_trans + i)*x.shape[1]:
                           (self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        pic[:, :, ::-1, :].reshape((self.get_data_dims(), x.shape[1])))

                    pic = np.rollaxis(pic, 0, 3)
                    # rotate +90
                    pic = np.rot90(pic)
                    target[:, (2*self.num_views/self.num_trans + i)*x.shape[1]:
                           (2*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                    # rotate +180
                    pic = np.rot90(pic)
                    target[:, (3*self.num_views/self.num_trans + i)*x.shape[1]:
                           (3*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                    # rotate +270
                    pic = np.rot90(pic)
                    target[:, (4*self.num_views/self.num_trans + i)*x.shape[1]:
                           (4*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                target[:, c] = pic.reshape((self.get_data_dims(),))

                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                else:
                    pic = np.rollaxis(pic, 0, 3)
                    pic = np.rot90(pic, nr.randint(1, 4))
                    pic = np.rollaxis(pic, 2, 0)
                target[:, x.shape[1]+c] = pic.reshape((self.get_data_dims(),))


class KaggleGalaxyZoo32DataProviderV2(LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range, init_epoch=1,
                 init_batchnum=None, dp_params={}, test=False):
        LabeledMemoryDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                           init_batchnum, dp_params, test)
        self.data_mean = self.batch_meta['data_mean']
        self.num_colors = 3
        self.img_size = 32
        self.data_mult = 8

        # Subtract the mean from the data and make sure that both data and
        # labels are in single-precision floating point.
        for d in self.data_dic:
            # This converts the data matrix to single precision and makes sure that it is C-ordered
            d['data'] = np.require((d['data'] - self.data_mean),
                                   dtype=np.single, requirements='C')
            d['labels'] = np.require(np.tile(d['labels'], (1, self.data_mult)),
                                     requirements='C')

        self.transformed_data = [np.zeros((self.get_data_dims(),
                                          self.data_dic[0]['data'].shape[1]*self.data_mult),
                                 dtype=np.single) for x in xrange(2)]
        self.batches_generated = 0

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledMemoryDataProvider.get_next_batch(self)
        transformed = self.transformed_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == transformed.shape[1]:
            self.batches_generated += 1
        else:
            transformed = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                                   dtype=np.single)

        self.__transform_data(datadic['data'], transformed)

        return epoch, batchnum, [transformed, datadic['labels']]

    # Returns the dimensionality of the two data matrices returned by get_next_batch
    # idx is the index of the matrix.
    def get_data_dims(self, idx=0):
        # A result has 37 fields
        return self.img_size**2 * self.num_colors if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.reshape(data.shape[1], 3, self.img_size,
                                                            self.img_size).swapaxes(1, 3)
                                                 .swapaxes(1, 2) / 255.0, dtype=np.single)

    def __transform_data(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])
        for c in xrange(x.shape[1]):    # loop over cases
            pic = y[:, :, :, c]
            target[:, c] = pic.reshape((self.get_data_dims(),))
            pic = np.rollaxis(pic, 0, 3)
            for i in xrange(1, 4):
                pic = np.rot90(pic)
                target[:, i*x.shape[1]+c] = (np.rollaxis(pic, 2, 0)
                                             .reshape((self.get_data_dims(),)))

            # horizontal reflection
            pic = y[:, :, ::-1, c]
            target[:, 4*x.shape[1]+c] = pic.reshape((self.get_data_dims(),))
            pic = np.rollaxis(pic, 0, 3)
            for i in xrange(5, 8):
                pic = np.rot90(pic)
                target[:, i*x.shape[1]+c] = (np.rollaxis(pic, 2, 0)
                                             .reshape((self.get_data_dims(),)))


class CroppedKaggleGalaxyZoo256DataProviderV2(LabeledDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                     init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 256 - self.border_size*2
        self.multiview = True
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.batch_meta['num_cases_per_batch']*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 256, 256))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledDataProvider.get_next_batch(self)
        datadic['data'] = np.require(datadic['data'], requirements='C')
        datadic['labels'] = np.require(np.tile(datadic['labels'], (1, self.data_mult)),
                                       requirements='C')

        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 256, 256, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for i in xrange(self.num_views/2):
                for c in xrange(x.shape[1]):    # loop over cases
                    startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                      nr.randint(0, self.border_size*2 + 1))
                    endY, endX = startY + self.inner_size, startX + self.inner_size
                    pic = y[:, startY:endY, startX:endX, c]
                    target[:, i*x.shape[1]+c] = pic.reshape((self.get_data_dims(),))
                    pic = pic[:, :, ::-1]
                    target[:, (i+self.num_views/2)*x.shape[1]+c] = (pic.reshape(
                                                                    (self.get_data_dims(),)))


class CroppedKaggleGalaxyZoo128DataProvider(LabeledDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                     init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 128 - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.batch_meta['num_cases_per_batch']*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 128, 128))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledDataProvider.get_next_batch(self)
        datadic['data'] = np.require(datadic['data'], requirements='C')
        datadic['labels'] = np.require(np.tile(datadic['labels'], (1, self.data_mult)),
                                       requirements='C')

        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 128, 128, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class CroppedKaggleGalaxyZoo32C1C6DataProvider(KaggleGalaxyZooC1C6LabeledMemoryDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        KaggleGalaxyZooC1C6LabeledMemoryDataProvider.__init__(self, data_dir, batch_range,
                                                              init_epoch, init_batchnum,
                                                              dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 32 - self.border_size*2
        self.multiview = test
        self.num_views = 5*2
        self.data_mult = self.num_views if self.multiview else 1
        self.num_colors = 3

        for d in self.data_dic:
            d['data'] = np.require(d['data'], requirements='C')
            d['label1'] = np.require(np.tile(d['label1'], (1, self.data_mult)),
                                     requirements='C')
            d['label6'] = np.require(np.tile(d['label6'], (1, self.data_mult)),
                                     requirements='C')
            d['label0'] = np.require(np.tile(d['label0'], (1, self.data_mult)),
                                     requirements='C')

        # cropped_data is a list with 2 elements: cropped_data[0] & cropped_data[1].
        # They have the same dimensionalities. So why do we need 2? See the answer right below
        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.data_dic[0]['data'].shape[1]*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 32, 32))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = KaggleGalaxyZooC1C6LabeledMemoryDataProvider.get_next_batch(self)

        # Here we cropped_data[0] and cropped_data[1] round-robin to avoid overlapping
        # assignments of 2 consecutive batches. Note that list assignment in Python is passing
        # reference instead of copying.
        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['label1'], datadic['label6'], datadic['label0']]

    def get_data_dims(self, idx=0):
        if idx == 0:
            return self.inner_size**2 * 3
        elif idx == 1:
            return 3
        elif idx == 2:
            return 2
        else:
            return 32

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 32, 32, x.shape[1])
        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/2):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    target[:, (self.num_views/2 + i)*x.shape[1]:
                           (self.num_views/2 + i + 1)*x.shape[1]] = (pic[:, :, ::-1, :].
                                                                     reshape((self.get_data_dims(),
                                                                              x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                target[:, :] = pic.reshape((self.get_data_dims(), x.shape[1]))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                if nr.randint(2) == 0:      # also flip the image with 50% probability
                    pic = pic[:, :, ::-1]
                target[:, c] = pic.reshape((self.get_data_dims(),))


class CroppedKaggleGalaxyZoo128DataProviderV2(LabeledDataProvider):
    def __init__(self, data_dir, batch_range=None, init_epoch=1,
                 init_batchnum=None, dp_params=None, test=False):
        LabeledDataProvider.__init__(self, data_dir, batch_range, init_epoch,
                                     init_batchnum, dp_params, test)

        self.border_size = dp_params['crop_border']
        self.inner_size = 128 - self.border_size*2
        self.multiview = dp_params['multiview_test'] and test
        self.num_trans = 5
        self.num_views = 1*5
        if test:
            self.data_mult = self.num_views
        else:
            self.data_mult = 1
        self.num_colors = 3

        self.cropped_data = [np.zeros((self.get_data_dims(),
                                      self.batch_meta['num_cases_per_batch']*self.data_mult),
                                      dtype=np.single) for x in xrange(2)]

        self.batches_generated = 0
        self.data_mean = (self.batch_meta['data_mean'].reshape((3, 128, 128))
                          [:, self.border_size:self.border_size+self.inner_size,
                           self.border_size:self.border_size+self.inner_size]
                          .reshape((self.get_data_dims(), 1)))

    def get_next_batch(self):
        epoch, batchnum, datadic = LabeledDataProvider.get_next_batch(self)
        datadic['data'] = np.require(datadic['data'], requirements='C')
        datadic['labels'] = np.require(np.tile(datadic['labels'], (1, self.data_mult)),
                                       requirements='C')

        cropped = self.cropped_data[self.batches_generated % 2]
        if datadic['data'].shape[1] == cropped.shape[1]:
            self.batches_generated += 1
        else:
            cropped = np.zeros((self.get_data_dims(), datadic['data'].shape[1]*self.data_mult),
                               dtype=np.single)

        self.__trim_borders(datadic['data'], cropped)
        cropped -= self.data_mean

        return epoch, batchnum, [cropped, datadic['labels']]

    def get_data_dims(self, idx=0):
        return self.inner_size**2 * 3 if idx == 0 else 37

    # Takes as input an array returned by get_next_batch
    # Returns a (numCases, imgSize, imgSize, 3) array which can be
    # fed to pylab for plotting.
    # This is used by shownet.py to plot test case predictions.
    def get_plottable_data(self, data):
        return np.require((data + self.data_mean).T.
                          reshape(data.shape[1], 3, self.inner_size, self.inner_size)
                          .swapaxes(1, 3).swapaxes(1, 2) / 255.0, dtype=np.single)

    def __trim_borders(self, x, target):
        y = x.reshape(3, 128, 128, x.shape[1])

        if self.test:   # don't need to loop over cases
            if self.multiview:
                start_positions = [(0, 0),  (0, self.border_size*2),
                                   (self.border_size, self.border_size),
                                   (self.border_size*2, 0),
                                   (self.border_size*2, self.border_size*2)]
                end_positions = [(sy+self.inner_size, sx+self.inner_size)
                                 for (sy, sx) in start_positions]
                for i in xrange(self.num_views/self.num_trans):
                    pic = y[:, start_positions[i][0]:end_positions[i][0],
                            start_positions[i][1]:end_positions[i][1], :]
                    target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                            x.shape[1]))
                    # horizontal reflection
                    target[:, (self.num_views/self.num_trans + i)*x.shape[1]:
                           (self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        pic[:, :, ::-1, :].reshape((self.get_data_dims(), x.shape[1])))

                    pic = np.rollaxis(pic, 0, 3)
                    # rotate +90
                    pic = np.rot90(pic)
                    target[:, (2*self.num_views/self.num_trans + i)*x.shape[1]:
                           (2*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                    # rotate +180
                    pic = np.rot90(pic)
                    target[:, (3*self.num_views/self.num_trans + i)*x.shape[1]:
                           (3*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                    # rotate +270
                    pic = np.rot90(pic)
                    target[:, (4*self.num_views/self.num_trans + i)*x.shape[1]:
                           (4*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                        np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
            else:
                # just take the center for now
                pic = y[:, self.border_size:self.border_size+self.inner_size,
                        self.border_size:self.border_size+self.inner_size, :]
                i = 0
                target[:, i*x.shape[1]:(i+1)*x.shape[1]] = pic.reshape((self.get_data_dims(),
                                                                       x.shape[1]))
                # horizontal reflection
                target[:, (self.num_views/self.num_trans + i)*x.shape[1]:
                       (self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                    pic[:, :, ::-1, :].reshape((self.get_data_dims(), x.shape[1])))

                pic = np.rollaxis(pic, 0, 3)
                # rotate +90
                pic = np.rot90(pic)
                target[:, (2*self.num_views/self.num_trans + i)*x.shape[1]:
                       (2*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                    np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                # rotate +180
                pic = np.rot90(pic)
                target[:, (3*self.num_views/self.num_trans + i)*x.shape[1]:
                       (3*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                    np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
                # rotate +270
                pic = np.rot90(pic)
                target[:, (4*self.num_views/self.num_trans + i)*x.shape[1]:
                       (4*self.num_views/self.num_trans + i + 1)*x.shape[1]] = (
                    np.rollaxis(pic, 2, 0).reshape((self.get_data_dims(), x.shape[1])))
        else:
            for c in xrange(x.shape[1]):    # loop over cases
                startY, startX = (nr.randint(0, self.border_size*2 + 1),
                                  nr.randint(0, self.border_size*2 + 1))
                endY, endX = startY + self.inner_size, startX + self.inner_size
                pic = y[:, startY:endY, startX:endX, c]
                tmp = nr.randint(0, 4)
                if tmp == 0:
                    pic = pic[:, :, ::-1]
                else:
                    pic = np.rollaxis(pic, 0, 3)
                    pic = np.rot90(pic, tmp)
                    pic = np.rollaxis(pic, 2, 0)
                target[:, c] = pic.reshape((self.get_data_dims(),))
